#include <ximera/stacktrace.hpp>
#include <ximera/common/utils.hpp>

#include <vector>

#if defined(XIMERA_WIN32)
#include <windows.h>
#include <DbgHelp.h>
#include <tlhelp32.h>
#endif

namespace ximera
{
    namespace detail
    {
#if defined(XIMERA_WIN32) && !defined(XIMERA_RELEASE)

        std::size_t const SYMBOL_NAME_MAXLEN = 1024;

        struct SymStartup
        {
            HANDLE process;
            SymStartup(HANDLE process)
                : process(process)
            {
                char current[MAX_PATH];
                std::string path;

                if (GetCurrentDirectoryA(MAX_PATH, current) > 0)
                {
                    path += current;
                    path += ";";
                }

                if (GetModuleFileNameA(NULL, current, MAX_PATH) > 0)
                {
                    std::string filePath = current;
                    std::string::size_type pos = filePath.find_last_of('\\');
                    if (pos != std::string::npos)
                        filePath.erase(pos);

                    path += filePath;
                    path += ";";
                }

                if (GetEnvironmentVariableA("_NT_SYMBOL_PATH", current, MAX_PATH) > 0)
                {
                    path += current;
                    path += ";";
                }

                if (GetEnvironmentVariableA("_NT_ALTERNATE_SYMBOL_PATH", current, MAX_PATH) > 0)
                {
                    path += current;
                    path += ";";
                }

                if (GetEnvironmentVariableA("SYSTEMROOT", current, MAX_PATH) > 0)
                {
                    path += current;
                    path += ";";
                    path += current;
                    path += "\\system32";
                    path += ";";
                }

                if (!SymInitialize(process, path.c_str(), FALSE))
                    throw 1;

                DWORD options = SymGetOptions();
                options |= SYMOPT_LOAD_LINES;
                options |= SYMOPT_FAIL_CRITICAL_ERRORS;

                options = SymSetOptions(options);
            }

            ~SymStartup()
            {
                if (process)
                    SymCleanup(process);
            }
        };

        class StackTraceImpl
        {
        private:            
            void load_modules(HANDLE process, DWORD processID)
            {
                HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, processID);
                if (snap == INVALID_HANDLE_VALUE)
                    return;

                MODULEENTRY32 entry;
                entry.dwSize = sizeof(entry);

                if (Module32First(snap, &entry))
                {
                    do
                    {
                        std::string fileName = string_util::wstr2str(entry.szExePath);
                        std::string moduleName = string_util::wstr2str(entry.szModule);
                        SymLoadModule64(process, NULL, fileName.c_str(), moduleName.c_str(), (DWORD64) entry.modBaseAddr, entry.modBaseSize);
                    }
                    while (Module32Next(snap, &entry));
                }

                CloseHandle(snap);
            }

            void retrieve_context(CONTEXT& context)
            {
                std::memset(&context, 0, sizeof(context));
                context.ContextFlags = CONTEXT_FULL;
                RtlCaptureContext(&context);
            }

            void retrieve_frame(CONTEXT& context, STACKFRAME64& frame, DWORD& imageType)
            {
                std::memset(&frame, 0, sizeof(frame));

            #ifdef _M_IX86
                imageType = IMAGE_FILE_MACHINE_I386;
                frame.AddrPC.Offset = context.Eip;
                frame.AddrPC.Mode = AddrModeFlat;
                frame.AddrFrame.Offset = context.Ebp;
                frame.AddrFrame.Mode = AddrModeFlat;
                frame.AddrStack.Offset = context.Esp;
                frame.AddrStack.Mode = AddrModeFlat;
            #elif _M_X64
                imageType = IMAGE_FILE_MACHINE_AMD64;
                frame.AddrPC.Offset = context.Rip;
                frame.AddrPC.Mode = AddrModeFlat;
                frame.AddrFrame.Offset = context.Rsp;
                frame.AddrFrame.Mode = AddrModeFlat;
                frame.AddrStack.Offset = context.Rsp;
                frame.AddrStack.Mode = AddrModeFlat;
            #elif _M_IA64
                imageType = IMAGE_FILE_MACHINE_IA64;
                frame.AddrPC.Offset = context.StIIP;
                frame.AddrPC.Mode = AddrModeFlat;
                frame.AddrFrame.Offset = context.IntSp;
                frame.AddrFrame.Mode = AddrModeFlat;
                frame.AddrBStore.Offset = context.RsBSP;
                frame.AddrBStore.Mode = AddrModeFlat;
                frame.AddrStack.Offset = context.IntSp;
                frame.AddrStack.Mode = AddrModeFlat;
            #else
                #error "Platform not supported!"
            #endif
            }

        public:
            void retrieve(StackTrace::ItemContainer& items)
            {
                HANDLE process = 0;

                try
                {
                    items.clear();

                    process = GetCurrentProcess();
                    SymStartup startup(process);

                    load_modules(process, GetCurrentProcessId());

                    HANDLE thread = GetCurrentThread();

                    CONTEXT context;
                    retrieve_context(context);

                    DWORD imageType = 0;
                    STACKFRAME64 frame;
                    retrieve_frame(context, frame, imageType);

                    std::vector<char> symbolData(sizeof(IMAGEHLP_SYMBOL64) + SYMBOL_NAME_MAXLEN, 0);
                    
                    IMAGEHLP_SYMBOL64* symbol = reinterpret_cast<IMAGEHLP_SYMBOL64*>(&symbolData[0]);                    
                    symbol->SizeOfStruct = sizeof(IMAGEHLP_SYMBOL64);
                    symbol->MaxNameLength = SYMBOL_NAME_MAXLEN;

                    IMAGEHLP_LINE64 line;
                    std::memset(&line, 0, sizeof(IMAGEHLP_LINE64));
                    line.SizeOfStruct = sizeof(IMAGEHLP_LINE64);

                    for (int frameNum = 0; true; ++frameNum)
                    {
                        if (!StackWalk64(imageType, process, thread, &frame, &context, NULL, &SymFunctionTableAccess64, &SymGetModuleBase64, NULL))
                            break;

                        if (frame.AddrPC.Offset == frame.AddrReturn.Offset)
                            break;                        

                        if (frame.AddrPC.Offset != 0)
                        {
                            StackItem item;

                            DWORD64 displacement64 = 0;
                            if (SymGetSymFromAddr64(process, frame.AddrPC.Offset, &displacement64, symbol))
                            {
                                char symbolName[SYMBOL_NAME_MAXLEN];
                                std::strncpy(symbolName, symbol->Name, SYMBOL_NAME_MAXLEN);

                                UnDecorateSymbolName(symbol->Name, symbolName, SYMBOL_NAME_MAXLEN, UNDNAME_COMPLETE);
                                item.name.assign(symbolName, symbolName + std::strlen(symbolName));
                            }

                            DWORD displacement = 0;
                            if (SymGetLineFromAddr64(process, frame.AddrPC.Offset, &displacement, &line))
                            {
                                item.line = line.LineNumber;
                                item.file.assign(line.FileName, line.FileName + std::strlen(line.FileName));
                            }

                            items.push_back(item);
                        }

                        if (frame.AddrReturn.Offset == 0)
                            break;
                    }
                }
                catch (...)
                {
                }                
            }
        };

#else
        class StackTraceImpl
        {
        public:
            void retrieve(StackTrace::ItemContainer& items)
            {
            }
        };
#endif
    }

    StackTrace::StackTrace()
        : _impl(new detail::StackTraceImpl)
    {
        _impl->retrieve(_items);

        if (_items.size() > 1)
            _items.erase(_items.begin(), _items.begin() + 2);
    }
}
